import fsm
import utils
import time
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine
from config_reader import DB
import models
from singleton import singleton
import threading
import logging
import const_define
import const
from config_reader import Config
from concurrent.futures import ThreadPoolExecutor


# create a thread pool using threads.ThreadPoolExecutor
@singleton
class ThreadPool(object):
    def __init__(self, max_workers):
        self.pool = ThreadPoolExecutor(max_workers)

    def submit(self, func, *args, **kwargs):
        return self.pool.submit(func, *args, **kwargs)

    def shutdown(self):
        self.pool.shutdown(wait=True)


@singleton
class Worker:
    def __init__(self):
        self.worker_id = utils.gen_worker_id()
        self.checkin_time = int(time.time())
        self.engine = create_engine(DB().engine_url)
        self.session = sessionmaker(bind=self.engine)()
        self.global_pusher = fsm.GlobalPusher()
        self.thread_pool = ThreadPool(max_workers=Config().max_workers)
        # register a new worker
        worker = models.Worker(
            worker_id=self.worker_id,
            checkin_time=self.checkin_time,
            shards_id=0,
            sharding_range=0
        )
        self.session.add(worker)
        self.session.commit()
        logging.info("[fsm: worker] worker {} registered".format(self.worker_id))
        # start a new thread for checking in
        threading.Thread(target=self.checkin, daemon=True).start()

    def leader_work(self):
        # the leader checks whether every worker is still alive, and re-assign the shards to every worker to guarantee
        # that every record will be push to an alive worker
        records_count = self.session.query(models.FsmRecord).count()
        alive_workers_count = 0
        worker_status = self.session.query(models.Worker).all()
        for record in worker_status:
            if int(time.time()) - record.checkin_time < const.MAX_CHECKIN_INTERVAL:
                alive_workers_count += 1

        shard_range = records_count / alive_workers_count

        # assign shards_id and sharding_range to every worker
        for index, record in enumerate(worker_status):
            if int(time.time()) - record.checkin_time < const.MAX_CHECKIN_INTERVAL:
                record.shards_id = int(index * shard_range)
                record.sharding_range = int(shard_range)
                self.session.commit()
                logging.info("[fsm: worker] worker {} assign shard {} to {}".format(self.worker_id, record.shards_id,
                                                                                    record.worker_id))

    def worker_work(self):
        # a worker initializes a global pusher and scans the records within its range
        record = self.session.query(models.Worker).filter(models.Worker.worker_id == self.worker_id).first()
        self.global_pusher.loop_work(record.shards_id, record.sharding_range)

    def check_identity(self):
        # check whether the worker is actually the leader
        record = self.session.query(models.Worker).filter(models.Worker.id == 1).first()
        if record.worker_id == 1:
            logging.info("[fsm: worker] worker {} is the leader".format(self.worker_id))
            self.leader_work()
        else:
            self.worker_work()

    def checkin(self):
        while True:
            self.checkin_time = int(time.time())
            self.session.query(models.Worker).filter(models.Worker.worker_id == self.worker_id) \
                .update({"checkin_time": self.checkin_time})
            logging.info("[fsm: worker] worker {} check in".format(self.worker_id))
            self.session.commit()
            # check whether the leader is still alive
            leader = self.session.query(models.Worker).filter(models.Worker.id == 1).first()
            if leader is None:
                leader_record = models.Worker(
                    id=1,
                    worker_id=self.worker_id,
                )
                self.session.add(leader_record)
                self.session.commit()
                logging.info("[fsm: worker] worker {} is the leader".format(self.worker_id))
            elif int(time.time()) - leader.checkin_time > const.MAX_CHECKIN_INTERVAL:
                # try to get the leader
                logging.info("[fsm: worker] worker {} try to get the leader".format(self.worker_id))
                # first, get the dispatch lock in mysql using FOR UPDATE
                self.session.query(models.Worker).filter(models.Worker.id == 1).with_for_update().first()
                # next, get the leader by changing the worker_id
                self.session.query(models.Worker).filter(models.Worker.id == 1).update({"worker_id": self.worker_id})
                # finally, commit the change
                self.session.commit()
                logging.info("[fsm: worker] worker {} get the leader".format(self.worker_id))

            # submit a task to thread pool to check the identity
            self.thread_pool.submit(self.check_identity)
            time.sleep(1)


if __name__ == '__main__':
    w = Worker()
    print(w.worker_id)
